#!/usr/bin/env python
import rospy
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
import cv2
import datetime
import threading
import queue

class ImageRecorder:
    def __init__(self):
        rospy.init_node('image_recorder', anonymous=True)
        self.bridge = CvBridge()
        self.video_out = None
        self.frame_count = 0
        self.queue = queue.Queue()
        self.lock = threading.Lock()
        self.running = True

        # 从参数服务器获取参数
        self.image_topic = rospy.get_param("~image_topic", "/camera/image_raw")
        self.out_fps = rospy.get_param("~output_fps", 30)
        self.out_width = rospy.get_param("~output_width", 640)
        self.out_height = rospy.get_param("~output_height", 480)
        self.out_filename = rospy.get_param("~output_filename", "output_video.mp4")

        # 获取当前日期和时间
        now = datetime.datetime.now()
        timestamp = now.strftime("%Y%m%d_%H%M%S")
        # 将时间戳添加到文件名
        self.out_filename = f"{self.out_filename.split('.')[0]}_{timestamp}.mp4"

        # 初始化视频输出，使用MP4编解码器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.video_out = cv2.VideoWriter(self.out_filename, fourcc, self.out_fps, (self.out_width, self.out_height))

        # 订阅图像话题
        rospy.Subscriber(self.image_topic, Image, self.image_callback)

        # 启动处理线程
        self.thread = threading.Thread(target=self.process_frames)
        self.thread.start()

    def image_callback(self, data):
        # 将ROS图像消息放入队列
        if self.running:
            self.queue.put(data)

    def process_frames(self):
        while self.running:
            try:
                # 从队列中获取图像
                data = self.queue.get(timeout=1)
                # 将ROS图像消息转换为OpenCV格式
                cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")

                # 调整图像大小以匹配输出视频的分辨率
                cv_image_resized = cv2.resize(cv_image, (self.out_width, self.out_height))

                # 写入视频帧
                self.video_out.write(cv_image_resized)
                self.frame_count += 1
                rospy.loginfo("Recorded frame %d", self.frame_count)
            except queue.Empty:
                continue

    def __del__(self):
        if self.video_out is not None:
            self.video_out.release()
            rospy.loginfo("Saved video to %s", self.out_filename)

if __name__ == '__main__':
    try:
        recorder = ImageRecorder()
        rospy.spin()
    except KeyboardInterrupt:
        rospy.loginfo("Shutting down, saving remaining frames...")
        recorder.running = False  # 停止处理线程
        recorder.thread.join()  # 等待处理线程结束
        if recorder.video_out is not None:
            recorder.video_out.release()
            rospy.loginfo("Saved video to %s", recorder.out_filename)
    except rospy.ROSInterruptException:
        pass
